#!/usr/bin/env python
print "importing libraries"
import rosbag
import cv_bridge

import sm
from sm import PlotCollection
import aslam_cv as acv
import aslam_cameras_april as acv_april
import aslam_cv_backend as acvb
import kalibr_common as kc
import kalibr_camera_calibration as kcc
import kalibr_rs_camera_calibration as rscc

import cv
import cv2
import os
import numpy as np
import pylab as pl
import argparse
import sys
import random
import signal

np.set_printoptions(suppress=True)

def __initBagDataset(bagfile, topic, from_to):
    print "\tDataset:          {0}".format(bagfile)
    print "\tTopic:            {0}".format(topic)
    reader = kc.BagImageDatasetReader(bagfile, topic, bag_from_to=from_to)
    print "\tNumber of images: {0}".format(reader.numImages())
    return reader

#available models
cameraModels = {
    # rolling shutter
    'pinhole-radtan-rs': acvb.DistortedPinholeRs,
    'pinhole-equi-rs':   acvb.EquidistantPinholeRs,
    'omni-radtan-rs':    acvb.DistortedOmniRs,
    # global shutter
    'pinhole-radtan':    acvb.DistortedPinhole,
    'pinhole-equi':      acvb.EquidistantPinhole,
    'omni-radtan':       acvb.DistortedOmni
}

def __signal_exit(signal, frame):
    sm.logWarn("Shutdown requested! (CTRL+C)")
    sys.exit(2)

def __parseArgs():
    class KalibrArgParser(argparse.ArgumentParser):
        def error(self, message):
            self.print_help()
            sm.logError('%s' % message)
            sys.exit(2)
        def format_help(self):
            formatter = self._get_formatter()
            formatter.add_text(self.description)
            formatter.add_usage(self.usage, self._actions,
                                self._mutually_exclusive_groups)
            for action_group in self._action_groups:
                formatter.start_section(action_group.title)
                formatter.add_text(action_group.description)
                formatter.add_arguments(action_group._group_actions)
                formatter.end_section()
            formatter.add_text(self.epilog)
            return formatter.format_help()

    usage = """
    Example usage to calibrate a camera system with a single rolling shutter camera using an aprilgrid.

    cam: pinhole model with radial-tangential distortion and rolling shutter for a single camera

    %(prog)s --model pinhole-radtan-rs --target aprilgrid.yaml \\
              --bag MYROSBAG.bag --topic /cam0/image_raw \\
              --feature-variance 1 --frame-rate 20

    example aprilgrid.yaml:
        target_type: 'aprilgrid'
        tagCols: 6
        tagRows: 6
        tagSize: 0.088  #m
        tagSpacing: 0.3 #as a fraction of tagSize"""

    parser = KalibrArgParser(description='Calibrate the intrinsics of a single rolling shutter camera.', usage=usage)
    parser.add_argument('--model', dest='model', help='The camera model to estimate. '
                        'Currently supported models are {0}, where the suffix \'-rs\' identifies rolling shutter models.'.format(', '.join(cameraModels.keys())), required=True)
    parser.add_argument('--frame-rate', dest='framerate', type=int, help='Approximate framerate of the camera.', required=True)

    groupSource = parser.add_argument_group('Data source')
    groupSource.add_argument('--bag', dest='bagfile', help='The bag file with the data.')
    groupSource.add_argument('--topic', dest='topic', help='The image topic.', required=True)
    groupSource.add_argument('--bag-from-to', metavar='bag_from_to', type=float, nargs=2, help='Use the bag data starting from up to this time [s].')

    groupTarget = parser.add_argument_group('Calibration target configuration')
    groupTarget.add_argument('--target', dest='targetYaml', help='Calibration target configuration as yaml file.', required=True)
    groupTarget.add_argument('--inverse-feature-variance', dest='inverseFeatureVariance', type=float, help='Estimated inverse variance of the feature detector.', required=True)

    groupOpt = parser.add_argument_group('Optimization options')
    groupOpt.add_argument('--max-iter', type=int, default=30, dest='max_iter', help='Max. iterations (default: %(default)s).', required=False)

    outputSettings = parser.add_argument_group('Output options')
    outputSettings.add_argument('--verbose', action='store_true', dest='verbose', help='Enable (really) verbose output (disables plots).')
    outputSettings.add_argument('--show-extraction', action='store_true', dest='showextraction', help='Show the calibration target extraction. (disables plots).')

    # print help if no argument is specified
    if len(sys.argv)==1:
        parser.print_help()
        sys.exit(2)

    # Parse the argument list
    try:
        parsed = parser.parse_args()
    except:
        sys.exit(2)

    # there is an issue with the gtk plot widget, so we cant plot if we have opencv windows open...
    # --> disable the plots in these special situations
    if parsed.showextraction or parsed.verbose:
        parsed.dontShowReport = True

    return parsed

# perform corner extraction from bagfile
# returns the observations
def __extractObservations(cameraGeometry, verbose=False):

    # extract the targets
    multithreading = False #not verbose
    observations = kc.extractCornersFromDataset(
        cameraGeometry.dataset,
        cameraGeometry.ctarget.detector,
        multithreading=multithreading,
        # transformation estimation will fail with rs cameras and significant distortions
        noTransformation=True
    )

    return observations

def main():
    parsed = __parseArgs()

    # logging modes
    if parsed.verbose:
        sm.setLoggingLevel(sm.LoggingLevel.Debug)
    else:
        sm.setLoggingLevel(sm.LoggingLevel.Info)

    # register signal handler
    signal.signal(signal.SIGINT, __signal_exit)

    ######
    ## load bagfile and extract targets:
    targetConfig = kc.CalibrationTargetParameters(parsed.targetYaml)
    dataset = __initBagDataset(parsed.bagfile, parsed.topic, parsed.bag_from_to)

    #create camera
    cameraModel = cameraModels[parsed.model]
    cameraGeometry = kcc.CameraGeometry(cameraModel, targetConfig, dataset, verbose=(parsed.verbose or parsed.showextraction))

    # extract observations
    observations = __extractObservations(cameraGeometry,(parsed.verbose or parsed.showextraction))

    # Calibration Configuration
    config = rscc.RsCalibratorConfiguration()
    config.inverseFeatureCovariance = 1.0/parsed.inverseFeatureVariance
    config.maxNumberOfIterations = parsed.max_iter
    config.framerate = parsed.framerate
    config.timeOffsetConstantSparsityPattern = 0.08
    config.timeOffsetPadding = 0.5

    ######
    ## Calibrate
    calibrator = rscc.RsCalibrator()
    calibrator.calibrate(cameraGeometry, observations, config)

if __name__ == "__main__":
    try:
        main()
    except Exception,e:
        sm.logError("Exception: {0}".format(e))
        sys.exit(-1)

